from numpy import *

def spike_detect(v_trace, threshold=0):
    ''' This method will detect the peaks of spikes that occur above a 
        threshold for a given voltage trace
        Params: v_trace = the voltage trace; threshold = min spike height
        Returns: An array of spike times'''
    spike_times = [];
    for v_i in range(len(v_trace)-1):
        #if the voltage trace crosses a threshold find the maximum over 1 msec
        if v_trace[v_i] <= threshold and v_trace[v_i+1] >= threshold:
            for search_area in range(v_i, len(v_trace)-1):
                if v_trace[search_area] > v_trace[search_area+1]:
                    spike_times.append(search_area)
                    break
    return spike_times

def burst_detect(spikes, intraburst_variance=1):
    ''' This method will find spike bursts in a voltage trace and return
        the burst start and end times
        Params: spikes = the spike times
                intraburst_variance = the maximum deviation from previous ISI
                    Value must be between 0 and 1
        Returns: An array of lists of start and end times
        '''
    bursts = []
    if len(spikes) < 10:
        print "There are not enough spikes to detect bursts"
        return 0
    else:
        isi_prev = spikes[1] - spikes[0] #find the first ISI
        burst_data = [spikes[0], spikes[1], 2] #the start of the first burst = first spike
        num_spikes = 2
        #for the rest of the spikes, find bursts
        for spike_i in range(1, len(spikes)-1):
            if num_spikes == 1:
                isi_prev = spikes[spike_i+1] - spikes[spike_i]
            isi_next = spikes[spike_i+1] - spikes[spike_i]
            ratio = float(isi_next)/float(isi_prev) #take the ratio of the ISIs
            isi_prev = isi_next
            #if the ratio differs by less than the allowed deviation, add that
            #spike to the burst and update the end time
            if ratio > (1-intraburst_variance) and ratio < (1+intraburst_variance):
                num_spikes = num_spikes + 1
                burst_data = [burst_data[0], spikes[spike_i+1], num_spikes]
            #if the next ISI is more than allowed variance longer than the previous, 
            #define as new burst
            else:
                bursts.append(burst_data)
                burst_data = [spikes[spike_i+1], spikes[spike_i+1], 1]
                num_spikes = 1
        bursts.append(burst_data)
    return bursts

def classify(v_trace, threshold=0, intraburst_variance=5, max_cv=0.2):
    ''' This function will take a voltage trace and classify the cell as either
        bursting, tonically active, non-periodic, or silent
        Params: v_trace = voltage trace; threshold = min spike threshold
        Returns: 0-silent; 1-tonically active; 2-bursting; 3-non-periodic
                 -1 - not enough spikes'''
    spikes = spike_detect(v_trace, threshold)
    if len(spikes) == 0: # if there are no spikes, return 0 (silent cell)
        return 0
    bursts = burst_detect(spikes, intraburst_variance) # find the bursts
    if bursts == 0: # if there are no bursts, return -1
        return -1
    interburst_intervals = []
    periods = []
    prev_burst_start = 0
    prev_burst_end = 0
    variance = 0
    for burst_i in range(len(bursts)): # for all the bursts, determine the periodicity
        if burst_i == 0: # if it's the first burst, initialize the values
            prev_burst_start = bursts[burst_i][0]
            prev_burst_end = bursts[burst_i][1]
        else:
            curr_burst_start = bursts[burst_i][0]
            curr_burst_end = bursts[burst_i][1]
            #find the interburst interval
            interburst_intervals.append(curr_burst_start - prev_burst_end)
            #find the period
            periods.append(curr_burst_start - prev_burst_start)
            prev_burst_start = curr_burst_start;
            prev_burst_end = curr_burst_end;
        if burst_i < len(bursts)-1:
            if bursts[burst_i][2] == bursts[burst_i+1][2]:
                #if two bursts have the same number of spikes, calculate the difference b/w corresponding ISI
                i_variance = interburst_variance(spikes,bursts[burst_i][0],bursts[burst_i][1],bursts[burst_i+1][0],bursts[burst_i+1][1])
                variance = abs(average(i_variance)-1.00)
        #if the number of spikes in the burst are greater than n/2, classify as tonically firing
        if bursts[burst_i][2] > len(spikes)/2:
            if len(bursts) < 3:
                return 1
    # find the coefficient of variation of the interburst periods
    mean_II = sum(interburst_intervals)/len(interburst_intervals)
    mean_period = sum(interburst_intervals)/len(interburst_intervals)
    from numpy import std
    std_II = std(interburst_intervals)
    std_period = std(periods)
    cv_II = std_II/mean_II
    cv_period = std_period/mean_period
    # if the CV is greater than the designated maximum, classify as non-periodic
    if cv_II > max_cv or cv_period > max_cv or variance > 0.01:
        return 3
    return 2
            
def interburst_variance(spikes,b1_on,b1_off,b2_on,b2_off):
    ''' This function compares the ISI in two bursts with an equal number of spikes
        Params: spikes = spike times; b1_on = time of first burst onset;
                b1_off = time of first burst offset; b2_on, b2_off
        Return: -1 if bursts are not same size
                an array of the ratios of the ISI
    '''
    burst1 = []
    burst2 = []
    # need to separate the spike times into the bursts
    for s_i in spikes:
        if s_i >= b1_on and s_i <= b1_off:
            burst1.append(s_i)
        elif s_i >= b2_on and s_i <= b2_off:
            burst2.append(s_i)
    # if bursts are not the same length, return -1
    if not len(burst1) == len(burst2):
        return -1
    isis1 = []
    isis2 = []
    variances = []
    # calculate the ISIs
    for s_i in range(len(burst1)-1):
        isis1.append(burst1[s_i+1]-burst1[s_i])
        isis2.append(burst2[s_i+1]-burst2[s_i])
    # get the ratios of the corresponding ISIs
    for i_index in range(len(isis1)):
        variances.append(float(isis1[i_index])/float(isis2[i_index]))
    return variances
